# Plotrix - hex grid plotting
# Ape - phylogenetic tree plotting
library(plotrix)
library(ape)


# overload color2D.matplot
color2D.matplot <- function (x, redrange = c(0, 1), greenrange = c(0, 1), bluerange = c(0, 
    1), extremes = NA, cellcolors = NA, show.legend = FALSE, 
    nslices = 10, xlab = "Column", ylab = "Row", do.hex = FALSE, 
    axes = TRUE, show.values = FALSE, vcol = "white", vcex = 1, 
    border = "black", na.color = NA, ...) 
{
    if (is.matrix(x) || is.data.frame(x)) {
        xdim <- dim(x)
        if (is.data.frame(x)) 
            x <- unlist(x)
        else x <- as.vector(x)
        oldpar <- par("xaxs", "yaxs", "xpd", "mar")
        par(xaxs = "i", yaxs = "i")
        if (do.hex) 
            #par(mar = c(5, 4, 4, 4))
        plot(c(0, xdim[2]), c(0, xdim[1]), xlab = xlab, ylab = ylab, 
            type = "n", axes = FALSE, ...)
        oldpar$usr <- par("usr")
        if (!do.hex) {
            box()
            pos <- 0
        }
        else pos <- -0.3
        if (axes) {
            xticks <- pretty(0:xdim[2])[-1]
            axis(1, at = xticks - 0.5, labels = xticks, pos = pos)
            yticks <- pretty(0:xdim[1])[-1]
            axis(2, at = xdim[1] - yticks + 0.5, yticks)
        }
        if (all(is.na(cellcolors))) 
            cellcolors <- color.scale(x, redrange, greenrange, 
                bluerange, extremes, na.color)
        if (do.hex) {
            par(xpd = TRUE)
            offset <- 0
            if (length(border) < xdim[1] * xdim[2]) 
                border <- rep(border, length.out = xdim[1] * 
                  xdim[2])
            for (row in 1:xdim[1]) {
                for (column in 0:(xdim[2] - 1)) {
                  hexagon(column + offset, xdim[1] - row, col = cellcolors[row + 
                    xdim[1] * column], border = border[row + 
                    xdim[1] * column])
                  if (show.values) 
                    text(column + offset + 0.5, xdim[1] - row + 
                      0.5, x[row + column * xdim[1]], col = vcol, 
                      cex = vcex)
                }
                offset <- ifelse(offset, 0, 0.5)
            }
            par(xpd = FALSE)
        }
        else {
            rect(sort(rep((1:xdim[2]) - 1, xdim[1])), rep(seq(xdim[1] - 
                1, 0, by = -1), xdim[2]), sort(rep(1:xdim[2], 
                xdim[1])), rep(seq(xdim[1], 1, by = -1), xdim[2]), 
                col = cellcolors, border = FALSE)
            if (show.values) 
                text(sort(rep((1:xdim[2]) - 0.5, xdim[1])), rep(seq(xdim[1] - 
                  0.5, 0, by = -1), xdim[2]), round(x, show.values), 
                  col = vcol, cex = vcex)
        }
        naxs <- which(is.na(x))
        xy <- par("usr")
        plot.din <- par("din")
        plot.pin <- par("pin")
        bottom.gap <- (xy[3] - xy[4]) * (plot.din[2] - plot.pin[2])/(2 * 
            plot.pin[2])
        grx1 <- xy[1]
        gry1 <- bottom.gap * 0.95
        grx2 <- xy[1] + (xy[2] - xy[1])/4
        gry2 <- bottom.gap * 0.8
        if (length(cellcolors) > 1) {
            colmat <- col2rgb(c(cellcolors[which.min(x)], cellcolors[which.max(x)]))
            redrange <- colmat[1, ]/255
            greenrange <- colmat[2, ]/255
            bluerange <- colmat[3, ]/255
        }
        rect.col <- color.scale(1:nslices, redrange, greenrange, 
            bluerange)
        if (show.legend) 
            color.legend(grx1, gry1, grx2, gry2, round(range(x, 
                na.rm = TRUE), 2), rect.col = rect.col)
        par(oldpar)
    }
    else cat("x must be a data frame or matrix\n")
}



# Colors used for number of selective hits
rainbowcol <- c("black", "violet", "blue", "darkgreen", "orange", "red")

for (sw in 709:1080){

crypts <- read.csv(paste("../sim",sw,"/sim",sw,".crypts",sep=""),header=FALSE)
# Read the patterns
patterns <- read.csv(paste("../sim",sw,"/sim",sw,".patterns",sep=""),header=FALSE)
phylogeny <- read.table(paste("../sim",sw,"/sim",sw,".phylogeny.log",sep=""),header=FALSE)

# Take phylo tree at end of simulation
tree <- read.tree(text=as.character(phylogeny[nrow(phylogeny),3]))
tips <- as.numeric(tree$tip.label)

# Do PCA on all unique neutral patterns
phylopatterns <- patterns[patterns[,1]%in%tips,]

# Take record of grid state at end of simulation
# .grid.log              date, generation, crypt_id, clone_id, pattern_id
glog <- read.csv(paste("../sim",sw,"/sim",sw,".grid.log",sep=""))
glog <- glog[glog[,1]==7300,]

# Assign colors by PCA to every pattern in the final patterns list
# PCA them
# col 1 is the name of the clone
# exclude selective loci, include only 100 neutral loci
pat <- phylopatterns[,2:101]
# PCA
pat_pca <- prcomp(pat)
# Make sure we have 3 principal componennts for red green blue
if(ncol(predict(pat_pca))>2){
red_col <- predict(pat_pca)[,1]
grn_col <- predict(pat_pca)[,2]
blu_col <- predict(pat_pca)[,3]
red_col <- red_col+abs(min(red_col))
red_col <- red_col/max(red_col)
grn_col <- grn_col+abs(min(grn_col))
grn_col <- grn_col/max(grn_col)
blu_col <- blu_col+abs(min(blu_col))
blu_col <- blu_col/max(blu_col)
if(!(is.nan(sum(red_col))|is.nan(sum(grn_col))|is.nan(sum(blu_col)))){
# Assign colors to edges
edge.neutral.colors <- NULL
edge.selective.colors <- NULL
j <- 1
for(i in 1:nrow(tree$edge)){
        if(tree$edge[i,2]==j){
		# j is an index 1..n where n is the total number of tips
		# tip 1 is clone_id 3, tip 2 is clone_id 300, and so on		
		ind <- which(phylopatterns[,1]==as.numeric(tree$tip.label[j]))
		edge.neutral.colors[i] <- rgb(red_col[ind],grn_col[ind],blu_col[ind])
		edge.selective.colors[i] <- rainbowcol[sum(phylopatterns[ind,107:111])+1]
		j <- j+1
        } else {
		ind <- which(phylopatterns[,1]==as.numeric(tree$tip.label[j]))
		edge.neutral.colors[i] <- rgb(red_col[ind],grn_col[ind],blu_col[ind])
		edge.selective.colors[i] <- rainbowcol[sum(phylopatterns[ind,107:111])+1]
        }	
}
pdf(paste("../sim",sw,"/sim",sw,".phylogeny.colored.by.neutral.patterns.pdf",sep=""),width=8.5,height=11)
par(mar=c(3,0,0,0.2))
plot(tree,type="c",x.lim=c(0,7300),no.margin=FALSE,root.edge=TRUE,show.node.label=TRUE,edge.color=edge.neutral.colors,show.tip.label=FALSE)
par(mgp=c(2,1,0))
axis(side=1,labels=c(0,5,10,15,20),at=c(0,1820,3640,5460,7300))
mtext(side=1,line=2,paste("Years, Sim",sw))
dev.off()

pdf(paste("../sim",sw,"/sim",sw,".phylogeny.colored.by.selective.patterns.pdf",sep=""),width=8.5,height=11)
par(mar=c(3,0,0,0.2))
plot(tree,type="c",x.lim=c(0,7300),no.margin=FALSE,root.edge=TRUE,show.node.label=TRUE,edge.color=edge.selective.colors,show.tip.label=FALSE)
par(mgp=c(2,1,0))
axis(side=1,labels=c(0,5,10,15,20),at=c(0,1820,3640,5460,7300))
mtext(side=1,line=2,paste("Years, Sim",sw))
legend(0,yjust=0,legend=paste("n=",c(0:5),sep=""),fill=rainbowcol)
dev.off()



# Loop all pattern_ids in glog vector and 
# get index of the first occurrence of clone_id in finalpatterns vector
ind <- match(glog[,4],phylopatterns[,1])

# last element of ind is length of ind = 90001
ind[length(ind)+1] <- length(ind)+1
ind[which(is.na(ind))] <- length(ind)
# the color of element 90001 is black = dead crypts
red_col[90001] <- 0
grn_col[90001] <- 0
blu_col[90001] <- 0
log.neutral.colors <- rgb(red_col[ind],grn_col[ind],blu_col[ind])

log.selective.colors <- apply(phylopatterns[,107:111],1,sum)+1
# black for dead crypts
log.selective.colors[90001] <- 1
log.selective.colors <- rainbowcol[log.selective.colors[ind]]

mat <- matrix(crypts[,1],nrow=300,ncol=300,byrow=TRUE)
mat.neutral.col <- mat
mat.selective.col <- mat
for(x in 1:300){
for(y in 1:300){
      mat.neutral.col[x,y]<-log.neutral.colors[(x-1)*300+y]	
      mat.selective.col[x,y]<-log.selective.colors[(x-1)*300+y]
}
}

pdf(paste("../sim",sw,"/sim",sw,".grid.colored.by.neutral.patterns.pdf",sep=""),width=3,height=3)
par(mar=c(0,0,0,0))
color2D.matplot(mat,cellcolors=mat.neutral.col,xlab=NA,ylab=NA,do.hex=TRUE,border=NA,axes=FALSE)
dev.off()

pdf(paste("../sim",sw,"/sim",sw,".grid.colored.by.selective.patterns.pdf",sep=""),width=3,height=3)
par(mar=c(0,0,0,0))
color2D.matplot(mat,cellcolors=mat.selective.col,xlab=NA,ylab=NA,do.hex=TRUE,border=NA,axes=FALSE)
dev.off()
}
}
print (paste("Sim",sw,"done."))
}




